"""
huggingface_embedding.py 模块
HuggingFace嵌入模型实现（支持自动检测设备、本地加载和自动下载）
"""
from typing import List
from pathlib import Path
import torch  # 新增：用于自动检测GPU
from sentence_transformers import SentenceTransformer
from src.embedding.base_embedding import BaseEmbedding

class HuggingFaceEmbedding(BaseEmbedding):
    def __init__(self, model_name: str, model_path: str, device: str = "auto"):
        self.model_name = model_name  # 模型名称（如"BAAI/bge-small-zh-v1.5"）
        self.model_path = model_path  # 本地存储根目录（如"./models/embeddings"）

        # 修复：自动解析设备参数（auto -> cuda/cpu）
        self.device = self._resolve_device(device)

        # 拼接完整本地路径（根目录 + 模型名，如"./models/embeddings/bge-small-zh-v1.5"）
        self.full_model_path = Path(model_path) / model_name.split("/")[-1]

        # 加载模型（优先本地，无则下载）
        self.model = self._load_model()

    def _resolve_device(self, device: str) -> str:
        """解析设备参数，将"auto"转换为实际设备（cuda/cpu）"""
        if device.lower() == "auto":
            # 自动检测：有GPU用cuda，否则用cpu
            return "cuda" if torch.cuda.is_available() else "cpu"
        # 支持显式设备（如"cpu"、"cuda:0"）
        return device

    def _load_model(self) -> SentenceTransformer:
        """加载模型（本地存在则直接加载，否则从HuggingFace下载并保存）"""
        try:
            # 检查本地模型是否存在且非空
            if self.full_model_path.exists() and any(self.full_model_path.glob("*")):
                print(f"✅ 从本地加载嵌入模型：{self.full_model_path}（设备：{self.device}）")
                return SentenceTransformer(
                    str(self.full_model_path),
                    device=self.device
                )
            else:
                # 本地无模型，从HuggingFace下载
                print(f"⚠️ 本地模型不存在，开始下载：{self.model_name}（设备：{self.device}）")
                print(f"📌 下载完成后将保存到：{self.full_model_path}")

                # 下载并加载模型
                model = SentenceTransformer(
                    self.model_name,
                    device=self.device
                )

                # 保存到本地（创建父目录，避免路径不存在错误）
                self.full_model_path.parent.mkdir(parents=True, exist_ok=True)
                model.save(str(self.full_model_path))
                print(f"✅ 模型已保存到本地：{self.full_model_path}")

                return model
        except Exception as e:
            # 详细错误信息，方便排查（如网络问题、路径权限等）
            raise RuntimeError(
                f"❌ 加载嵌入模型失败：{str(e)}\n"
                f"请检查：1. 模型路径是否正确 {self.full_model_path}\n"
                f"       2. 网络是否通畅（如需下载）\n"
                f"       3. 设备参数是否有效（当前：{self.device}）"
            )

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """为文档列表生成嵌入向量"""
        return self.model.encode(texts, show_progress_bar=False).tolist()

    def embed_query(self, text: str) -> List[float]:
        """为查询文本生成嵌入向量"""
        return self.model.encode(text, show_progress_bar=False).tolist()